#!/usr/bin/env python3
# encoding: utf-8
# @data:2023/07/13
# @author:aiden
# 巡线+上下阶梯+跨栏x2
import rospy
import math
import signal
from ainex_sdk import pid, misc, common
from ainex_kinematics.gait_manager import GaitManager
from ainex_example.color_common import Common
from ainex_example.pid_track import PIDTrack
from ainex_example.visual_patrol_match import VisualPatrolMatch
from ainex_example.approach_object import ApproachObject
from ainex_interfaces.srv import SetString
from ainex_interfaces.msg import ObjectsInfo, ColorDetect, ROI

class CombinationNode(Common):
    # 按顺序检测三个roi，如果检测到黑线立刻跳出
    # y_min, y_max, x_min, x_max分别表示占图像的比例, 即实际大小为y_min*height
    line_roi = [(5 / 12, 6 / 12, 0, 1),
                (6 / 12, 7 / 12, 0, 1),
                (7 / 12, 8 / 12, 0, 1)
                ]

    stairs_roi = [0, 1, 0.2, 0.8]
    hurdles_roi = [1 / 5, 1, 0, 1]
    ball_roi = [0, 1, 0, 1] 
    # 所需动作的名称
    climb_stairs_action_name = 'climb_stairs_competition_fast0'
    descend_stairs_action_name = 'descend_stairs_fast1'
    back = 'back_step_fast'
    go = 'forward_one_step_fast'

    # 图像处理时缩放到这个分辨率， 不建议修改
    image_process_size = [160, 120]

    enter_climb_stairs_y = 110 / 480  # 当检测到的标识像素坐标y值占图像的比例大于此值时进入此阶段，进入上台阶，减慢速度
    climb_stairs_y_stop = 100 / 480  # 当检测到的标识像素坐标y值占图像的比例在此值附近(范围可在ApproachObject里设置)时停止横向移动


    descend_stairs_y_stop = 320 / 480 #下台阶的距离


    def __init__(self, name):
        rospy.init_node(name)
        self.name = name
        self.rl_dis = None
        self.gait_manager = GaitManager()
        self.go_gait_param = self.gait_manager.get_gait_param()
        self.go_gait_param['body_height'] = 0.025
        self.go_gait_param['step_height'] = 0.02
        self.go_gait_param['hip_pitch_offset'] = 15
        self.go_gait_param['z_swap_amplitude'] = 0.006
        self.go_dsp = [400, 0.2, 0.023]
        self.go_arm_swap = 30
        self.count = 0
        self.running = True
        self.objects_info = []

        self.current_state = "walk"      #当前状态
        #状态列表
        self.state = {'walk':   [[500, 260], ['white', self.line_roi, self.image_process_size, self.set_visual_patrol_color], False],
                      'up_yellow_climb_stairs':   [[500, 260], ['yellow', self.stairs_roi, self.image_process_size, self.set_stairs_color], False],
                      'up_climb_stairs':   [[500, 260], ['blue', self.stairs_roi, self.image_process_size, self.set_stairs_color], False],
                      'down_descend_stairs':   [[500, 260], ['blue', self.stairs_roi, self.image_process_size, self.set_stairs_color], False],
                      'ball': [[500, 450], ['blue', self.ball_roi, self.image_process_size, self.set_ball_color], False]}
        
        self.head_pan_init = self.state[self.current_state][0][0]   # 左右舵机的初始值
        self.head_tilt_init = self.state[self.current_state][0][1]  # 上下舵机的初始值
        super().__init__(name, self.head_pan_init, self.head_tilt_init)

        self.calib_config = common.get_yaml_data('/home/ubuntu/ros_ws/src/ainex_example/config/calib.yaml')

        self.approach_object = ApproachObject(self.gait_manager)
        self.visual_patrol_match = VisualPatrolMatch(self.gait_manager)
        
        # 头部的pid追踪
        self.head_pan_range = [125, 875]  # 左右转动限制在这个范围， 125为右
        self.head_tilt_range = [260, 500]  # 上下限制在这个范围， 250为下
        self.pid_rl = pid.PID(0.1, 0.0, 0.001)
        self.pid_ud = pid.PID(0.1, 0.0, 0.001)
        self.head_pan_init = 500  # 左右舵机的初始值
        self.head_tilt_init = 300 # 上下舵机的初始值
        self.rl_track = PIDTrack(self.pid_rl, self.head_pan_range, self.head_pan_init)
        self.ud_track = PIDTrack(self.pid_ud, self.head_tilt_range, self.head_tilt_init)

  # 躯体的追踪参数
        self.yaw_stop = 40  # 躯体停止转动时头部左右舵机脉宽值和中位500的差值

        signal.signal(signal.SIGINT, self.shutdown)
        
        # 订阅颜色识别结果
        rospy.Subscriber('/object/pixel_coords', ObjectsInfo, self.get_color_callback)
        rospy.Service('~set_color', SetString, self.set_color_srv_callback)  # 设置颜色
        self.motion_manager.run_action('walk_ready')

        if rospy.get_param('~start', True):
            # 通知颜色识别准备，此时只显示摄像头原画
            self.enter_func(None)
            self.start_srv_callback(None)
            common.loginfo('start ombination')

    def shutdown(self, signum, frame):
        self.running = False
        common.loginfo('%s shutdown' % self.name)

    def set_visual_patrol_color(self, color, roi, image_process_size):
        # 设置巡线颜色
        line_param = ColorDetect()
        line_param.color_name = color
        line_param.use_name = True
        line_param.detect_type = 'line'
        line_param.image_process_size = image_process_size
        line_param.line_roi.up.y_min = int(roi[0][0] * image_process_size[1])
        line_param.line_roi.up.y_max = int(roi[0][1] * image_process_size[1])
        line_param.line_roi.up.x_min = int(roi[0][2] * image_process_size[0])
        line_param.line_roi.up.x_max = int(roi[0][3] * image_process_size[0])

        line_param.line_roi.center.y_min = int(roi[1][0] * image_process_size[1])
        line_param.line_roi.center.y_max = int(roi[1][1] * image_process_size[1])
        line_param.line_roi.center.x_min = int(roi[1][2] * image_process_size[0])
        line_param.line_roi.center.x_max = int(roi[1][3] * image_process_size[0])

        line_param.line_roi.down.y_min = int(roi[2][0] * image_process_size[1])
        line_param.line_roi.down.y_max = int(roi[2][1] * image_process_size[1])
        line_param.line_roi.down.x_min = int(roi[2][2] * image_process_size[0])
        line_param.line_roi.down.x_max = int(roi[2][3] * image_process_size[0])

        line_param.min_area = 1
        line_param.max_area = image_process_size[0] * image_process_size[1]
        
        return line_param

    def set_stairs_color(self, color, roi, image_process_size):
        # 设置台阶标志颜色
        stairs_param = ColorDetect()
        stairs_param.color_name = color
        stairs_param.detect_type = 'side'
        stairs_param.use_name = True
        stairs_param.image_process_size = image_process_size
        stairs_param.roi.y_min = int(roi[0] * image_process_size[1])
        stairs_param.roi.y_max = int(roi[1] * image_process_size[1])
        stairs_param.roi.x_min = int(roi[2] * image_process_size[0])
        stairs_param.roi.x_max = int(roi[3] * image_process_size[0])
        stairs_param.min_area = 10*20
        stairs_param.max_area = image_process_size[0]*image_process_size[1]
        
        return stairs_param

    def set_ball_color(self, color, roi, image_process_size):
        # 设置追踪颜色
        ball_param = ColorDetect()
        ball_param.color_name = color
        ball_param.detect_type = 'circle'
        ball_param.use_name = True
        ball_param.image_process_size = self.image_process_size
        ball_param.roi.y_min = int(roi[0] * image_process_size[1])
        ball_param.roi.y_max = int(roi[1] * image_process_size[1])
        ball_param.roi.x_min = int(roi[2] * image_process_size[0])
        ball_param.roi.x_max = int(roi[3] * image_process_size[0])
        ball_param.min_area = 20
        ball_param.max_area = self.image_process_size[0]*self.image_process_size[1]

        return ball_param

    def set_color_srv_callback(self, msg):
        # 设置颜色
        stairs_param = self.set_stairs_color(self.state['climb_stairs'][1][0])
        line_param = self.set_visual_patrol_color(self.state['visual_patrol'][1][0])
        ball_param = self.set_ball_color(self.state['ball'][1][0])

        self.detect_pub.publish([line_param, stairs_param, ball_param])
        common.loginfo('%s set_color' % self.name)
        
        return [True, 'set_color']

    def get_color_callback(self, msg):
        # 获取颜色识别结果
        self.objects_info = msg.data
        for object_info in self.objects_info:
            if object_info.type == 'circle':
                object_info.x = object_info.x - self.calib_config['center_x_offset']
                self.rl_dis, self.ud_dis = self.head_track_process(object_info)
                if self.rl_dis is not None:
                    self.body_track_process(self.rl_dis, self.ud_dis, object_info)
                    self.rl_dis = None
                    self.count_miss = 0
                    self.start_find_ball = False

    def state_init(self, current_state):
        # 不同阶段的初始化
        if self.state[current_state][2] == False:
            self.state[current_state][2] = True
            self.init_action(self.state[current_state][0][0], self.state[current_state][0][1])  # 头部姿态
            param1 = self.state[current_state][1][3](self.state[current_state][1][0], self.state[current_state][1][1], self.state[current_state][1][2])
            self.detect_pub.publish([param1])  # 颜色检测设置
            common.loginfo(current_state + ' init')

    def state_init_walk(self, current_state):
        # 不同阶段的初始化
        if self.state[current_state][2] == False:
            self.state[current_state][2] = True
            self.init_action(self.state[current_state][0][0], self.state[current_state][0][1])  # 头部姿态
            param1 = self.state[current_state][1][3](self.state[current_state][1][0], self.state[current_state][1][1], self.state[current_state][1][2])
            param2 = self.state['up_yellow_climb_stairs'][1][3](self.state['up_yellow_climb_stairs'][1][0], self.state['up_yellow_climb_stairs'][1][1], self.state['up_yellow_climb_stairs'][1][2])
            self.detect_pub.publish([param1, param2])  # 颜色检测设置
            common.loginfo(current_state + ' init')
     
    #进入上台阶
    def enter_climb_stairs(self, stairs_data):
        if stairs_data is not None:
            print('max(stairs_data.y, stairs_data.left_point[1], stairs_data.right_point[1])',max(stairs_data.y, stairs_data.left_point[1], stairs_data.right_point[1]))
            print('self.enter_climb_stairs_y * stairs_data.height',self.enter_climb_stairs_y * stairs_data.height)
            if min(stairs_data.y, stairs_data.left_point[1], stairs_data.right_point[1]) > self.enter_climb_stairs_y * stairs_data.height:
                return True
            else:
                self.count = 0
        return False
    
    def exit_climb_stairs(self, stairs_data):
        # 上阶梯处理
        if stairs_data is None:  
            self.gait_manager.set_step(self.go_dsp, 0.005, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
        if stairs_data is not None:
            print("stairs_data.y:",stairs_data.y)
            print("self.climb_stairs_y_stop*stairs_data.height:",self.climb_stairs_y_stop*stairs_data.height)
            print('stairs_data.angle',stairs_data.angle)
            if abs(stairs_data.angle) >= 15:
                yaw = math.copysign(5, 0 - stairs_data.angle)
                self.gait_manager.set_step(self.go_dsp, 0, 0, yaw, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            elif stairs_data.y > self.climb_stairs_y_stop*stairs_data.height:
                self.gait_manager.disable()  # 关闭步态控制
                common.loginfo('climb_stairs_competition')
                self.motion_manager.run_action(self.climb_stairs_action_name)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.go)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.climb_stairs_action_name)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.go)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.climb_stairs_action_name)
                rospy.sleep(0.8)
                return True
            else:
                self.gait_manager.set_step(self.go_dsp, 0.005, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                
        return False

    def exit_descend_stairs(self, stairs_data):
        # 下阶梯处理
        print('stairs_data.y', stairs_data.y)
        if stairs_data is None:  
            self.gait_manager.set_step(self.go_dsp, 0.005, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
        if stairs_data is not None:

            if abs(stairs_data.angle) >= 10:
                yaw = math.copysign(5, 0 - stairs_data.angle )
                self.gait_manager.set_step(self.go_dsp, 0, 0, yaw, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
            elif stairs_data.y > self.descend_stairs_y_stop*stairs_data.height:
                self.gait_manager.disable()  # 关闭步态控制
                common.loginfo('climb_stairs')
                self.motion_manager.run_action(self.descend_stairs_action_name)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.back)
                self.motion_manager.run_action(self.back)
                self.motion_manager.run_action(self.back)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.descend_stairs_action_name)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.back)
                self.motion_manager.run_action(self.back)
                self.motion_manager.run_action(self.back)
                rospy.sleep(0.8)
                self.motion_manager.run_action(self.descend_stairs_action_name)
                rospy.sleep(0.8)
                return True
            else:
                self.gait_manager.set_step(self.go_dsp, 0.005, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)

        return False


    # 头部PID跟踪
    def head_track_process(self, object_info):
        # 头部追踪
        if abs(object_info.x - object_info.width/2) < 10:
            object_info.x = object_info.width/2
        if abs(object_info.y - object_info.height/2) < 10:
            object_info.y = object_info.height/2
        rl_dis = self.rl_track.track(object_info.x, object_info.width/2)
        ud_dis = self.ud_track.track(object_info.y, object_info.height/2)
        self.motion_manager.set_servos_position(20, [[23, int(rl_dis)], [24, int(ud_dis)]])
        return rl_dis, ud_dis
    
    # 躯体追踪
    def body_track_process(self, rl_dis, ud_dis, ball_data):                                                
        # 左右根据头部左右舵机位置值进行调整
        if ball_data is not None:
            if abs(rl_dis - 500) < 20:
                yaw_output =0
            else:
                yaw_output = math.copysign(5, rl_dis - 500)
            self.gait_manager.set_step(self.go_dsp, 0.02, 0, yaw_output, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)


    def run(self):
        ball_init = True
        while self.running:
            if self.start:
                # 获取识别结果
                line_data = None
                side_data = None
                ball_data = None
                for object_info in self.objects_info:
                    if object_info.type == 'line':
                        line_data = object_info
                    if object_info.type == 'side':
                        side_data = object_info
                
                # 当前阶段处理完成
                if self.current_state == 'walk':
                    if line_data is not None and side_data is None:
                        self.visual_patrol_match.process(line_data.x, line_data.width)
                        line_data = None
                    elif line_data is None and side_data is None:
                        self.gait_manager.set_step(self.go_dsp, 0.02, 0, 0, self.go_gait_param, arm_swap=self.go_arm_swap, step_num=0)
                    if self.enter_climb_stairs(side_data):
                        print(1)
                        self.current_state = 'up_climb_stairs'
                        self.state[self.current_state][2] = False  # 重新初始化当前阶段
                        side_data = None
                        common.loginfo('exit yellow_climb_stairs ---> enter while_climb_stairs')
                    else:
                        rospy.sleep(0.8)
                
                elif self.current_state == 'up_climb_stairs':
                    if self.exit_climb_stairs(side_data):
                        self.current_state = 'down_descend_stairs'
                        self.state[self.current_state][2] = False  # 重新初始化当前阶段
                        common.loginfo('exit yellow_climb_stairs ---> enter while_climb_stairs')
                    else:
                        rospy.sleep(0.8)
                elif self.current_state == 'down_descend_stairs':
                    if self.exit_descend_stairs(side_data):
                        print(4)
                        self.current_state = 'ball'
                        self.state[self.current_state][2] = False
                        common.loginfo('exit red_descend_stairs ---> enter while_descend_stairs')
                    else:
                        rospy.sleep(0.8)

                if self.current_state == 'walk':
                    self.state_init_walk(self.current_state)
                elif self.current_state == 'ball' and ball_init:
                    self.state_init(self.current_state)
                    ball_init = False
                else:
                    self.state_init(self.current_state)

                rospy.sleep(0.01)  # 防止空载
            else:
                rospy.sleep(0.01)
        
        self.init_action(self.head_pan_init, self.head_tilt_init)
        self.stop_srv_callback(None)
        rospy.signal_shutdown('shutdown')

if __name__ == "__main__":
    CombinationNode('footstep').run()
